#include <iostream>
#include <string>
#include <sstream>
#include <cstring>

#include "ntfsdump.h"
#include "fstypes.h"

using namespace std;

namespace ntfs {

CDatadump::CDatadump(int cols):
    mColInRow(cols)
{

}

void CDatadump::print(unsigned char* buf, int bufLen)
{
    int i;
    string str;

    cout.width(10);
    cout.fill(' ');
    cout << "";


    for (i=0; i < mColInRow; i++)
    {
        if ( (i % mColInRow != 0) && (i % 0x10 == 0) )
            cout << "  ";

        cout.width(2);
        cout.fill('0');
        cout << hex << (i) << " ";
    }

    cout << "\n\n";

    for (i=0; i< bufLen; i++)
    {
        if (i % mColInRow == 0)
        {
            cout << "0x";
            cout.width(4);
            cout.fill('0');
            cout << hex << ((i / mColInRow) * mColInRow) << "    ";
        }

        if ( (i % mColInRow != 0) && (i % 0x10 == 0) )
            cout << "  ";

        cout.width(2);
        cout.fill('0');

        int b = (unsigned char)buf[i];
        cout << hex << b << " ";


        if ( (b >= 32) && (b < 0xff))
            str += buf[i];
        else
            str += '.';


        if ((i+1) % mColInRow == 0)
        {
            cout << "        "  << str;
            cout << "\n";
            str.clear();
            if (i==511)
                cout << "\n";
        }
    }
}

void dumpBootBlock(PBOOT_BLOCK pBootBlock)
{
    cout << "BOOT BLOCK (" << dec << sizeof(BOOT_BLOCK) << "):\n";
    cout << "Disk geometry: \n";
    cout << "BytesPerSector: " << pBootBlock->BytesPerSector << "\n";
    cout << "SectorsPerCluster: " << (int)pBootBlock->SectorsPerCluster << "\n";
    cout << "ClustersPerFileRecord: " << pBootBlock->ClustersPerFileRecord << "(0x" << hex << pBootBlock->ClustersPerFileRecord << dec << ")" << "\n";

    int BytesPerFileRecord = pBootBlock->ClustersPerFileRecord < 0x80
    ? pBootBlock->ClustersPerFileRecord * pBootBlock->SectorsPerCluster * pBootBlock->BytesPerSector
    : 1 << (0x100 - pBootBlock->ClustersPerFileRecord);

    cout << "BytesPerFileRecord: " << BytesPerFileRecord << "\n";
    cout << "ClustersPerIndexBlock: " << pBootBlock->ClustersPerIndexBlock << "\n";
    cout << "SectorsPerTrack: " << pBootBlock->SectorsPerTrack << "\n";
    cout << "NumberOfHeads: " << pBootBlock->NumberOfHeads << "\n";
    cout << "BootSectors: " << pBootBlock->BootSectors << "\n";
    cout << "PartitionOffset: " << pBootBlock->PartitionOffset << "\n";
    cout << "TotalSectors: " << pBootBlock->TotalSectors << "\n";


    cout << "\nMft: \n";
    cout << "MftStartLcn: " << pBootBlock->MftStartLcn << "\n";
    cout << "Mft2StartLcn: " << pBootBlock->Mft2StartLcn << "\n";

    UINT64 mftOffset = pBootBlock->BytesPerSector * pBootBlock->SectorsPerCluster * pBootBlock->MftStartLcn;
    cout << "mftOffset: " << mftOffset << "\n";

    cout << "\nVolume: \n";
    cout << "VolumeSerialNumbe: " << pBootBlock->VolumeSerialNumber << "\n";
    cout << "MediaType: " << pBootBlock->MediaType << "\n";

    cout << "\nMzb: \n";
    cout << "Mbz1: " << (int)pBootBlock->Mbz1 << "\n";
    cout << "Mbz2: " << pBootBlock->Mbz2 << "\n";
    cout << "Mbz3: " << pBootBlock->Mbz3 << "\n";

    char czFormat[9] = {0};
    memcpy(czFormat, pBootBlock->Format, 8);
    cout << "Format: " << czFormat << "\n";
    cout << "BootSignature: " <<  pBootBlock->BootSignature << "\n";
}

template <class T>
void printInt(string sTitle, T val, int iIndent)
{
    string sIndent;
    for (int i=0; i < iIndent; i++)
        sIndent += "    ";
    cout << sIndent << sTitle << ": " << dec << (INT64)val << hex << " (0x" << (INT64)val << ")\n";
}

void printBool(string sTitle, bool val, int iIndent)
{
    string sIndent;
    for (int i=0; i < iIndent; i++)
        sIndent += "    ";
    cout << sIndent << sTitle << ": " << boolalpha << val << "\n";
}

void printString(string sTitle, string val, int iIndent)
{
    string sIndent;
    for (int i=0; i < iIndent; i++)
        sIndent += "    ";
    cout << sIndent << sTitle << ": " << val << "\n";
}

void printRawData(string sTitle, PUINT8 pData, int cbData, int iIndent)
{
    string sIndent;
    for (int i=0; i < iIndent; i++)
        sIndent += "    ";
    UINT8* cData = new UINT8[cbData+1];
    memcpy(cData, pData, cbData);
    cData[cbData] = 0;
    cout << sIndent << sTitle << ": " << cData << "\n";
    delete[] cData;
}

void dumpFileRec(CFileRecord& pFileRec, bool WithAttributes)
{
    cout << "FILE_RECORD_HEADER:\n";
    cout << "    Ntfs (NTFS_RECORD_HEADER):\n";
    printRawData("Type", PUINT8(&(pFileRec->Ntfs.Type)), 4, 2);
    printInt("UsaOffset", pFileRec->Ntfs.UsaOffset, 2);
    printInt("UsaCount", pFileRec->Ntfs.UsaCount, 2);
    printInt("LogSN", pFileRec->Ntfs.LogSN, 2);

    printInt("SequenceNumber", pFileRec->SequenceNumber, 1);
    printInt("LinkCount", pFileRec->LinkCount, 1);
    printInt("AttributesOffset", pFileRec->AttributesOffset, 1);
    printInt("Flags", pFileRec->Flags, 1);
    printInt("BytesInUse", pFileRec->BytesInUse, 1);
    printInt("BytesAllocated", pFileRec->BytesAllocated, 1);
    printInt<UINT64>("BaseFileRecord", pFileRec->BaseFileRecord & 0x0000FFFFFFFFFFFF, 1);
    printInt("NextAttributeNumber", pFileRec->NextAttributeNumber, 1);
    printInt("IndexNumber", pFileRec->IndexNumber, 1);
    cout << endl;   
    if (WithAttributes)
        for (unsigned int i = 0; i < pFileRec.getAttributesCount(); i++)
            dumpAttribute(*pFileRec.getAttribute(i));

}

void dumpUSNArray(UINT16* pUSN, UINT16 count)
{
    printInt("USN", pUSN[0], 0);
    cout << "USArray[" << count-1 <<"]: ";
    for (int i = 1; i < count; i++)
        cout << pUSN[i] << " ";
}

string getAttrTypeName(ATTRIBUTE_TYPE attrType)
{
    switch (attrType)
    {
        case AttributeStandardInformation:  return "AttributeStandardInformation";
        case AttributeAttributeList:  return "AttributeAttributeList";
        case AttributeFileName:  return "AttributeFileName";
        case AttributeObjectId:  return "AttributeObjectId";
        case AttributeSecurityDescriptor:  return "AttributeSecurityDescriptor";
        case AttributeVolumeName:  return "AttributeVolumeName";
        case AttributeVolumeInformation:  return "AttributeVolumeInformation";
        case AttributeData:  return "AttributeData";
        case AttributeIndexRoot:  return "AttributeIndexRoot";
        case AttributeIndexAllocation:  return "AttributeIndexAllocation";
        case AttributeBitmap:  return "AttributeBitmap";
        case AttributeReparsePoint:  return "AttributeReparsePoint";
        case AttributeEAInformation:  return "AttributeEAInformation";
        case AttributeEA:  return "AttributeEA";
        case AttributePropertySet:  return "AttributePropertySet";
        case AttributeLoggedUtilityStream:  return "AttributeLoggedUtilityStream";
    }
    return "";
}

void dumpAttributeResident(CResidentAttribute* pAttr)
{
    CResidentAttribute resAttr = *pAttr;
    printInt("ValueLength", resAttr->ValueLength, 2);
    printInt("ValueOffset", resAttr->ValueOffset, 2);
    printInt("Flags", resAttr->Flags, 2);
}

void dumpAttributeNonResident(CNonResidentAttribute* pAttr)
{
    CNonResidentAttribute nonResAttr = *pAttr;
    printInt("LowVcn", nonResAttr->LowVcn, 2);
    printInt("HighVcn", nonResAttr->HighVcn, 2);
    printInt("RunArrayOffset", nonResAttr->RunArrayOffset, 2);
    printInt("CompressionUnit", nonResAttr->CompressionUnit, 2);
    printInt("AllocatedSize", nonResAttr->AllocatedSize, 2);
    printInt("DataSize", nonResAttr->DataSize, 2);
    printInt("InitializedSize", nonResAttr->InitializedSize, 2);
    printInt("CompressedSize", nonResAttr->CompressedSize, 2);
    cout << "    Runs:\n" << endl;
    for (unsigned int i= 0; i < nonResAttr.mRun.mRunList.size(); i++)
    {
        cout << "["
             << dec << nonResAttr.mRun.mRunList[i].LowVCN << ": "
             << nonResAttr.mRun.mRunList[i].LowLCN << " - "
             << nonResAttr.mRun.mRunList[i].ClustersCount << "]" << endl;
    }
}

wstring getAttributeName(const CAttribute& Attribute)
{
    if (Attribute->NameLength == 0)
        return L"-";
    wchar_t* ptrName = (wchar_t*)IncPTR(PCATTRIBUTE(Attribute), Attribute->NameOffset);
    return wstring(ptrName, Attribute->NameLength);
}

void dumpAttribute(CAttribute& pAttribute)
{
    if (pAttribute->Nonresident)
        cout << "ATTRIBUTE (Nonresident):\n";
    else
        cout << "ATTRIBUTE (Resident):\n";
    printString("AttributeType", getAttrTypeName(pAttribute->AttributeType), 1);
    printInt("Length", pAttribute->Length, 1);
    printBool("NonResident", pAttribute->Nonresident, 1);
    printInt("NameLength", pAttribute->NameLength, 1);
    printInt("NameOffset", pAttribute->NameOffset, 1);
    printInt("Flags", pAttribute->Flags, 1);
    printInt("AttributeNumber", pAttribute->AttributeNumber, 1);
    wcout << L"    Name: " << getAttributeName(pAttribute) << endl;
    if (pAttribute->Nonresident)
        dumpAttributeNonResident(dynamic_cast<CNonResidentAttribute*>(&pAttribute));
    else
        dumpAttributeResident(dynamic_cast<CResidentAttribute*>(&pAttribute));
    cout << endl;
}

wstring getAttributeListName(const PATTRIBUTE_LIST pAttribute)
{
    if (pAttribute->NameLength == 0)
        return L"-";
    wchar_t* ptrName = (wchar_t*)IncPTR(PCATTRIBUTE(pAttribute), pAttribute->NameOffset);
    return wstring(ptrName, pAttribute->NameLength);
}

void dumpAttributeList(PATTRIBUTE_LIST attrList)
{
    attrList->FileReferenceNumber = attrList->FileReferenceNumber & 0x0000FFFFFFFFFFFF;
    printString("AttributeType", getAttrTypeName(attrList->AttributeType), 1);
    printInt("Length", attrList->Length, 1);
    printInt("NameLength", attrList->NameLength, 1);
    printInt("NameOffset", attrList->NameOffset, 1);
    printInt("LowVcn", attrList->LowVcn, 1);
    printInt("FileReferenceNumber", attrList->FileReferenceNumber, 1);
    printInt("AttributeNumber", attrList->AttributeNumber, 1);
    wcout << L"Name: " << getAttributeListName(attrList) << endl;
}

} //namespace ntfs

